/*
 * PsFileStream.cpp
 *
 *  Created on: 2018年7月8日
 *      Author: zhengboyuan
 */

#include "PsFileStream.h"


#ifndef MAKEWORD
typedef unsigned char BYTE;
typedef unsigned short WORD;
#define MAKEWORD(a,b) ((WORD) (((BYTE) (a)) | ((WORD) ((BYTE) (b))) << 8))
#endif //

static size_t   SYNC_BYTES_LENGTH = 4;  /// including start code

static size_t   MIN_HEADER_SIZE = 6;

static size_t   MAX_PS_LENGTH = 0x100000;
static size_t   MAX_PES_LENGTH = 0xFFFF;
static size_t   MAX_ES_LENGTH = 0x100000;



bool PsFileStream::getPesData(PsPacket& pkt, PsPacket& out)
{
    bool found = false;
    if (pkt.type == kPsAudioPacket || pkt.type == kPsVideoPacket)
    {
        size_t hdrSize = sizeof(pes_header_t) + sizeof(optional_pes_header_t);

        pes_header_t* header = (pes_header_t*)pkt.data;
        unsigned short pes_packet_length = MAKEWORD(header->pes_packet_length[1], header->pes_packet_length[0]);

        optional_pes_header_t* pes_header = (optional_pes_header_t*)(pkt.data + sizeof(pes_header_t));
        hdrSize += pes_header->PES_header_data_length;

        out.pos = pkt.pos + hdrSize;
        out.data = pkt.data + hdrSize;
        out.size = pkt.size - hdrSize;

        found = true;
    }
    return found;
}

bool PsFileStream::getTime(PsPacket& pkt, time_t& t)
{
    if (pkt.type != kPsPackHeader)
    {
        return false;
    }

    ps_header* header = (ps_header*)pkt.data;
    if (header->pack_stuffing_length >= 6)
    {
        uint8_t* data = pkt.data + sizeof(ps_header);
        int index = 0;
        uint8_t year = data[index ++];
        uint8_t month = data[index ++];
        uint8_t day = data[index ++];
        uint8_t hour = data[index ++];
        uint8_t minute = data[index ++];
        uint8_t second = data[index ++];

        struct tm tmNow;
        memset(&tmNow, 0, sizeof(tmNow));
        tmNow.tm_year = year + 2000 - 1900;
        tmNow.tm_mon = month - 1;
        tmNow.tm_mday = day;
        tmNow.tm_hour = hour;
        tmNow.tm_min = minute;
        tmNow.tm_sec = second;

        t = mktime(&tmNow);
    }

    return true;
}

bool PsFileStream::getTimeStamp(PsPacket& pkt, int64_t& pts, int64_t& dts)
{
    bool found = false;
    if (pkt.type == kPsAudioPacket || pkt.type == kPsVideoPacket)
    {
        size_t hdrSize = sizeof(pes_header_t) + sizeof(optional_pes_header_t);

        pes_header_t* header = (pes_header_t*)pkt.data;
        unsigned short pes_packet_length = MAKEWORD(header->pes_packet_length[1], header->pes_packet_length[0]);

        optional_pes_header_t* pes_header = (optional_pes_header_t*)(pkt.data + sizeof(pes_header_t));
        hdrSize += pes_header->PES_header_data_length;

        bool hasPts = false;

        if ((pes_header->PES_header_data_length >= 5) && (pes_header->PTS_DTS_flags != 0))
        {
            /// parse ts
            int index = 9;
            uint8_t* data = pkt.data;

            int64_t b30 = (data[index] & 0x0E) >> 1;
            index++;
            int64_t b22 = data[index];
            index++;
            int64_t b15 = data[index] >> 1;
            index++;
            int64_t b8 = data[index];
            index++;
            int64_t b0 = data[index] >> 1;
            pts = b0 + (b8 << 7) + (b15 << 15) + (b22 << 22) + (b30 << 30);

            hasPts = true;
        }

        found = hasPts;
    }
    return found;
}

bool PsFileStream::getTimeStamp(PsPacket& pkt, int64_t& pts)
{
	int64_t dts = 0;
	return getTimeStamp(pkt, pts, dts);
}


PsFileStream::PsFileStream():
        m_state(kStart),
        m_pos(),
        m_curPacket(),
        m_curTime(),
        m_lastPacket(),
        m_lastTime()
{
    m_buffer.ensure(MAX_PES_LENGTH);
}

PsFileStream::~PsFileStream()
{
    close();
}

bool PsFileStream::open(const char* filepath)
{
    m_ifs.open(filepath, std::ios::binary);
    if (!m_ifs.is_open())
    {
        return false;
    }

    m_filepath = filepath;

    return true;
}

void PsFileStream::close()
{
    if (m_ifs.is_open())
    {
        m_ifs.close();
    }
}

bool PsFileStream::isOpen()
{
    return m_ifs.is_open();
}

int64_t PsFileStream::tellg()
{
    return m_ifs.tellg();
}

bool PsFileStream::seekg(int64_t pos)
{
	m_ifs.clear();
    return !m_ifs.seekg(pos).eof();
}

bool PsFileStream::eof()
{
    return m_ifs.eof();
}

const std::string& PsFileStream::getFilePath() const
{
    return m_filepath;
}

bool PsFileStream::read(PsPacket& pkt)
{
    bool found = false;
    while (!m_ifs.eof() && !found)
    {
        int ch = m_ifs.get();
        if (ch == -1)
        {
            break;
        }
        else
        {
            found = handleByte((uint8_t)ch, pkt);

            if (found)
            {
                m_lastPacket = m_curPacket;
                m_lastTime = m_curTime;

                m_curPacket.type = pkt.type;
                m_curPacket.pos = pkt.pos;
                PsFileStream::getTime(pkt, m_curTime);
            }
        }
    }
    return found;
}

bool PsFileStream::handleByte(uint8_t ch, PsPacket& pkt)
{
    bool found = false;
    if (m_state == kStart)
    {
        if (ch == 0x00)
        {
            m_state = kSingleZero;

			m_pos = m_ifs.tellg();
			m_pos = m_pos - 1;
        }
    }
    else if (m_state == kSingleZero)
    {
        if (ch == 0x00)
        {
            m_state = kDoubleZero;
        }
        else
        {
            m_state = kStart;
            m_pos = 0;
        }
    }
    else if (m_state == kDoubleZero)
    {
        if (ch == 0x00)
        {
            m_pos += 1;
            m_state = kDoubleZero;
        }
        else if (ch == 0x01)
        {
            m_state = kStartCode;
        }
        else
        {
            m_state = kStart;
            m_pos = 0;
        }
    }
    else if (m_state == kStartCode)
    {
        if (ch == 0xBA)
        {
            m_state = kPackHeader;

            m_buffer.clear();

            uint8_t startCode[] = {0, 0, 1, ch};
            m_buffer.write(startCode, sizeof(startCode));
            readToBuffer(m_buffer, sizeof(ps_header) - 4);

            ps_header* header = (ps_header*)m_buffer.data();
            if (header->pack_stuffing_length > 0)
            {
                readToBuffer(m_buffer, header->pack_stuffing_length);
            }

            buildPacket(pkt, kPsPackHeader);

            m_state = kStart;
            m_pos = 0;

            found = true;
        }
        else if (ch == 0xBB)
        {
            m_state = kSystemHeader;

            m_buffer.clear();

            uint8_t startCode[] = {0, 0, 1, ch};
            m_buffer.write(startCode, sizeof(startCode));
            readToBuffer(m_buffer, sizeof(sh_header) - 4);

            sh_header* header = (sh_header*)m_buffer.data();
            unsigned short header_length = MAKEWORD(header->header_length[1], header->header_length[0]);
            size_t wanted = SYNC_BYTES_LENGTH + 2 + header_length;
			wanted -= m_buffer.length();
            readToBuffer(m_buffer, wanted);

            buildPacket(pkt, kPsSystemHeader);

            m_state = kStart;
            m_pos = 0;

            found = true;
        }
        else if (ch == 0xBC)
        {
            m_state = kPsmHeader;

            m_buffer.clear();

            uint8_t startCode[] = {0, 0, 1, ch};
            m_buffer.write(startCode, sizeof(startCode));
            readToBuffer(m_buffer, sizeof(psm_header_t) - 4);

            psm_header_t* header = (psm_header_t*)m_buffer.data();
            uint16_t stream_map_length = MAKEWORD(header->program_stream_map_length[1], header->program_stream_map_length[0]);
            uint16_t stream_info_length = MAKEWORD(header->program_stream_info_length[1], header->program_stream_info_length[0]);
            uint16_t es_map_length = MAKEWORD(header->elementary_stream_map_length[1], header->elementary_stream_map_length[0]);
            size_t wanted = SYNC_BYTES_LENGTH + 2 + stream_map_length;
			wanted -= m_buffer.length();
            readToBuffer(m_buffer, wanted);

            buildPacket(pkt, kPsPsmHeader);

            m_state = kStart;
            m_pos = 0;

            found = true;
        }
        else if (ch >= 0xC0 && ch <= 0xDF)
        {
            m_state = kPesAudioPacket;

            m_buffer.clear();

            uint8_t startCode[] = {0, 0, 1, ch};
            m_buffer.write(startCode, sizeof(startCode));
            readToBuffer(m_buffer, sizeof(pes_header_t) - 4);

            pes_header_t* header = (pes_header_t*)m_buffer.data();
            unsigned short pes_packet_length = MAKEWORD(header->pes_packet_length[1], header->pes_packet_length[0]);
            size_t wanted = SYNC_BYTES_LENGTH + 2 + pes_packet_length;
			wanted -= m_buffer.length();
            readToBuffer(m_buffer, wanted);

            buildPacket(pkt, kPsAudioPacket);

            m_state = kStart;
            m_pos = 0;

            found = true;
        }
        else if (ch >= 0xE0 && ch <= 0xEF)
        {
            m_state = kPesVideoPacket;

            m_buffer.clear();

            uint8_t startCode[] = {0, 0, 1, ch};
            m_buffer.write(startCode, sizeof(startCode));
            readToBuffer(m_buffer, sizeof(pes_header_t) - 4);

            pes_header_t* header = (pes_header_t*)m_buffer.data();
            unsigned short pes_packet_length = MAKEWORD(header->pes_packet_length[1], header->pes_packet_length[0]);
            size_t wanted = SYNC_BYTES_LENGTH + 2 + pes_packet_length;
			wanted -= m_buffer.length();
            readToBuffer(m_buffer, wanted);

            buildPacket(pkt, kPsVideoPacket);

            m_state = kStart;
            m_pos = 0;

            found = true;
        }
        else if (ch >= 0xBD && ch <= 0xBF)
        {
            m_state = kStart;
            m_pos = 0;
        }
        else if (ch >= 0xF0 && ch <= 0xF8)
        {
            m_state = kStart;
            m_pos = 0;
        }
        else
        {
            m_state = kStart;
            m_pos = 0;
        }
    }
    return found;
}

int PsFileStream::readToBuffer(comn::ByteBuffer& buffer, int size)
{
    if (size <= 0)
    {
        return 0;
    }

    buffer.ensure(buffer.length() + size);

    uint8_t* data = buffer.data();
    int length = buffer.length();
    data += length;

    if (m_ifs.read((char*)data, size))
    {
        buffer.resize(length + size);
        return size;
    }
    return 0;
}

void PsFileStream::buildPacket(PsPacket& pkt, PsPacketType type)
{
    pkt.pos = m_pos;
    pkt.type = type;
    pkt.data = m_buffer.data();
    pkt.size = m_buffer.length();
}

time_t PsFileStream::getLastTime()
{
    return m_lastTime;
}

PsPacket PsFileStream::getLastPacket()
{
    return m_lastPacket;
}

time_t PsFileStream::getCurTime()
{
    return m_curTime;
}

PsPacket PsFileStream::getCurPacket()
{
    return m_curPacket;
}

bool PsFileStream::read(int size, comn::ByteBuffer& buffer)
{
	int length = readToBuffer(buffer, size);
	return (length == size);
}
